#!/usr/bin/env python3
"""
CogView MCP HTTP Server

An HTTP server that provides MCP-compatible endpoints for image generation using CogView
with optional background removal functionality.
"""

import asyncio
import json
import logging
from typing import Any, Dict
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import uvicorn

from .cogview_client import CogViewClient
from .background_remover import BackgroundRemover
from .config import Config

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(
    title="CogView MCP HTTP Server",
    description="HTTP server providing MCP-compatible endpoints for CogView image generation",
    version="1.0.0"
)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Initialize clients
cogview_client = CogViewClient()
bg_remover = BackgroundRemover()


# Request models
class GenerateImageRequest(BaseModel):
    prompt: str
    model: str = "cogview-3-flash"
    size: str = "1024x1024"


class GenerateImageTransparentRequest(BaseModel):
    prompt: str
    model: str = "cogview-3-flash"
    size: str = "1024x1024"
    background_removal_model: str = "u2net"


class MCPToolRequest(BaseModel):
    name: str
    arguments: Dict[str, Any]


@app.get("/")
async def root():
    """Root endpoint with server information."""
    return {
        "name": "CogView MCP HTTP Server",
        "version": "1.0.0",
        "description": "HTTP server providing MCP-compatible endpoints for CogView image generation",
        "endpoints": {
            "/tools": "List available MCP tools",
            "/tools/call": "Call MCP tools",
            "/generate/image": "Generate image directly",
            "/generate/image/transparent": "Generate transparent background image directly"
        }
    }


@app.get("/tools")
async def list_tools():
    """List available MCP tools."""
    return {
        "tools": [
            {
                "name": "generate_image",
                "description": "Generate an image using CogView based on a text prompt",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "prompt": {
                            "type": "string",
                            "description": "Text prompt describing the image to generate"
                        },
                        "model": {
                            "type": "string",
                            "description": "CogView model to use",
                            "default": "cogview-3-flash",
                            "enum": ["cogview-3-flash"]
                        },
                        "size": {
                            "type": "string",
                            "description": "Image size in format 'widthxheight'",
                            "default": "1024x1024",
                            "enum": [
                                "1024x1024", "768x1024", "1024x768", "1152x896", 
                                "896x1152", "1216x832", "832x1216", "1344x768", 
                                "768x1344", "1536x640", "640x1536"
                            ]
                        }
                    },
                    "required": ["prompt"]
                }
            },
            {
                "name": "generate_image_transparent",
                "description": "Generate an image with transparent background by removing the background from a CogView generated image",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "prompt": {
                            "type": "string",
                            "description": "Text prompt describing the image to generate"
                        },
                        "model": {
                            "type": "string",
                            "description": "CogView model to use",
                            "default": "cogview-3-flash",
                            "enum": ["cogview-3-flash"]
                        },
                        "size": {
                            "type": "string",
                            "description": "Image size in format 'widthxheight'",
                            "default": "1024x1024",
                            "enum": [
                                "1024x1024", "768x1024", "1024x768", "1152x896", 
                                "896x1152", "1216x832", "832x1216", "1344x768", 
                                "768x1344", "1536x640", "640x1536"
                            ]
                        },
                        "background_removal_model": {
                            "type": "string",
                            "description": "Background removal model to use",
                            "default": "u2net",
                            "enum": ["u2net", "u2netp", "u2net_human_seg", "silueta"]
                        }
                    },
                    "required": ["prompt"]
                }
            }
        ]
    }


@app.post("/tools/call")
async def call_tool(request: MCPToolRequest):
    """Call MCP tools."""
    try:
        if request.name == "generate_image":
            return await generate_image_tool(request.arguments)
        elif request.name == "generate_image_transparent":
            return await generate_image_transparent_tool(request.arguments)
        else:
            raise HTTPException(status_code=400, detail=f"Unknown tool: {request.name}")
    except Exception as e:
        logger.error(f"Error in tool {request.name}: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/generate/image")
async def generate_image_direct(request: GenerateImageRequest):
    """Generate image directly via REST API."""
    return await generate_image_tool(request.dict())


@app.post("/generate/image/transparent")
async def generate_image_transparent_direct(request: GenerateImageTransparentRequest):
    """Generate transparent background image directly via REST API."""
    return await generate_image_transparent_tool(request.dict())


async def generate_image_tool(arguments: Dict[str, Any]):
    """Generate an image using CogView."""
    prompt = arguments.get("prompt", "")
    model = arguments.get("model", "cogview-3-flash")
    size = arguments.get("size", Config.DEFAULT_IMAGE_SIZE)
    
    if not prompt:
        raise HTTPException(status_code=400, detail="prompt is required")
    
    logger.info(f"Generating image with prompt: {prompt}, model: {model}, size: {size}")
    
    try:
        # Generate image using CogView
        image_url = await cogview_client.generate_image(prompt, model, size)
        
        # Download and encode image
        image_data = await cogview_client.download_image(image_url)
        
        return {
            "result": [
                {
                    "type": "text",
                    "text": f"Image generated successfully with prompt: '{prompt}'"
                },
                {
                    "type": "image",
                    "data": image_data,
                    "mimeType": "image/png"
                }
            ]
        }
        
    except Exception as e:
        logger.error(f"Error generating image: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error generating image: {str(e)}")


async def generate_image_transparent_tool(arguments: Dict[str, Any]):
    """Generate an image with transparent background."""
    prompt = arguments.get("prompt", "")
    model = arguments.get("model", "cogview-3-flash")
    size = arguments.get("size", Config.DEFAULT_IMAGE_SIZE)
    bg_model = arguments.get("background_removal_model", "u2net")
    
    if not prompt:
        raise HTTPException(status_code=400, detail="prompt is required")
    
    logger.info(f"Generating transparent image with prompt: {prompt}, model: {model}, size: {size}")
    
    try:
        # Generate image using CogView
        image_url = await cogview_client.generate_image(prompt, model, size)
        
        # Download image
        image_data = await cogview_client.download_image(image_url)
        
        # Remove background
        transparent_image_data = await bg_remover.remove_background(image_data, bg_model)
        
        return {
            "result": [
                {
                    "type": "text",
                    "text": f"Transparent image generated successfully with prompt: '{prompt}'"
                },
                {
                    "type": "image",
                    "data": transparent_image_data,
                    "mimeType": "image/png"
                }
            ]
        }
        
    except Exception as e:
        logger.error(f"Error generating transparent image: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error generating transparent image: {str(e)}")


def start_server(host: str = "localhost", port: int = 8000):
    """Start the HTTP server."""
    logger.info(f"Starting CogView MCP HTTP Server on {host}:{port}")
    uvicorn.run(app, host=host, port=port)


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="CogView MCP HTTP Server")
    parser.add_argument("--host", default="localhost", help="Host to bind to")
    parser.add_argument("--port", type=int, default=8000, help="Port to bind to")
    
    args = parser.parse_args()
    start_server(args.host, args.port)
